from util import clean_dir
from workflow import NovelWorkflow, AgentState
from nodes import NovelNodes
import re

def parse_chapters(outline: str) -> list:
    """从大纲中解析章节列表"""
    chapter_pattern = r'第(\\d+)章\\s+.*'
    return re.findall(chapter_pattern, outline)

REF = """
# 要素
- 校园甜文
- 爽文
- 金手指：花钱系统
"""

def main():
    # 初始化工作流
    workflow = NovelWorkflow()
    app = workflow.compile_workflow()
    
    # 初始状态配置
    initial_state = AgentState(
        ref=REF,
        outline="",
        refine_count=0,
        current_chapter=1,  # 初始化为默认章节编号0
        chapter_content="",
        word_count=0
    )
    
    clean_dir('output')

    # 生成并迭代大纲
    state = app.invoke(initial_state, config={'recursion_limit': 15})
    

    
    # 解析章节并循环生成
    total_chapters = len(re.findall(r'第\\d+章', state['outline']))
    for chapter_num in range(1, total_chapters+1):
        chapter_state = AgentState(
            outline=state['outline'],
            refine_count=state['refine_count'],
            current_chapter=chapter_num,
            chapter_content="",
            word_count=0
        )
        result = app.invoke(chapter_state)
        
        if not writeFile(f'output/chapter_{chapter_num}.txt', result['chapter_content']):
            raise Exception(f'章节{chapter_num}文件生成失败')
        
        print(f'已生成第{chapter_num}章: 字数{result["word_count"]}')

if __name__ == "__main__":
    main()